
from matplotlib.pyplot import close, plot
import pandas as pd
import backtrader as bt
from backtrader_plotting import Bokeh
from backtrader_plotting.schemes import Tradimo
from my_plot import *
from py_mysql import *
import numpy as np
import talib as ta
import math

class PandasData_more(bt.feeds.PandasData):
    # pass
    lines = ('Bmd1','Bmd2','Bdn1','Bdn2',) # 要添加的线
    # 设置 line 在数据源上的列位置
    params=(
        ('Bmd1', -1),
        ('Bmd2', -1),
        ('Bdn1', -1),
        ('Bdn2', -1),
           ) 
    # -1表示自动按列明匹配数据，也可以设置为线在数据源中列的位置索引 (('pe',6),('pb',7),)

# 创建策略继承bt.strategy
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod', 15),
        # 判断是否输出该日志
        ('printlog', False),
    )

    # 初始化部分指标

    def __init__(self):
        # 下单数量
        self.volSize = 10
        # 信号状态值
        self.signalStatus = None
        # 持仓情况
        self.positionStatus = None
        # 止损价格
        self.stopPrice = None
        # 成交数量
        self.dealNum = 0
        # 盈利单数
        self.profitSize = 0
        # 亏损单数
        self.lossSize = 0

    # 策略核心代码
    def next(self):
        # 当前日期
        self.time = str(bt.num2date(self.lines.datetime[0]))
        self.positionSize = self.position.size

        buySignal = self.datas[0].Bdn1[-2] > self.datas[0].Bdn2[-2] and self.datas[0].close[-2] > self.datas[0].Bmd1[-2]
        sellSignal = self.datas[0].Bdn1[-2] < self.datas[0].Bdn2[-2] and self.datas[0].close[-2] < self.datas[0].Bmd1[-2]

        if buySignal:
            self.signalStatus = 1
        elif sellSignal:
            self.signalStatus = -1

        stopSize = 27 * 1 * self.volSize
        # 空仓
        if not self.position:
            if self.signalStatus == 1:
                self.buy(size = self.volSize)
                self.positionStatus = 'buy'
                self.stopPrice = self.datas[0].close - stopSize
            elif self.signalStatus == -1:
                self.sell(size = self.volSize)
                self.positionStatus = 'sell'
                self.stopPrice = self.datas[0].close + stopSize
        # 有持仓
        else:
            if self.positionStatus == 'sell' and self.datas[0].close > self.stopPrice:
                print('做空达到止损出场')
                self.close()
                self.positionStatus = None
                return
            elif self.positionStatus == 'buy' and self.datas[0].close < self.stopPrice:
                print('做多达到止损出场')
                self.close()
                self.positionStatus = None
                return

            if self.signalStatus == 1 and self.positionStatus == 'sell':
                self.close()
                self.buy(size = self.volSize)
                self.positionStatus = 'buy'
                self.stopPrice = self.datas[0].close - stopSize
            elif self.signalStatus == -1 and self.positionStatus == 'buy':
                self.close()
                self.sell(size = self.volSize)
                self.positionStatus = 'sell'
                self.stopPrice = self.datas[0].close + stopSize
            
    # 输出分位数
    def quantile_p(self,data, p):
        pos = (len(data) + 1)*p
        #pos = 1 + (len(data)-1)*p
        pos_integer = int(math.modf(pos)[1])
        pos_decimal = pos - pos_integer
        Q = data[pos_integer - 1] + (data[pos_integer] - data[pos_integer - 1])*pos_decimal
        return Q

    # 日志输出
    def log(self, txt, dt=None, doprint=False):
        if self.params.printlog or doprint:
            # 记录策略的执行日志
            dt = dt or self.datas[0].datetime.date(0)
            print(f'{dt.isoformat()},{txt}')

    # 订单状态通知，买入卖出都是下单
    def notify_order(self, order):
        # 提交了/接受了,  买/卖订单什么都不做
        if order.status in [order.Submitted, order.Accepted]:
            return
        # 检查一个订单是否完成
        # 注意:当资金不足时，broker会拒绝订单
        if order.status in [order.Completed]:
            self.dealNum += 1
            if order.isbuy():
                print('日期：{}---买入价格:{}---买入手续费{}'.format(self.time,order.executed.price,
                      order.executed.comm))
                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            elif order.issell():
                print('日期：{}---卖出价格:{}---卖出手续费{}'.format(self.time,order.executed.price,
                      order.executed.comm))
            self.bar_executed = len(self)
        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('订单取消/保证金不足/拒绝', doprint=False)
        # 其他状态记录为：无法挂单
        self.order = None

    # 交易状态通知，一买一卖算交易（交易净利润）
    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        if trade.pnlcomm > 0:
            self.profitSize += 1
        else:
            self.lossSize += 1

    # 策略结束时，多用于参数调优
    def stop(self):
        print('总成交笔数：{}'.format(self.dealNum))
        print('盈利笔数：{}'.format(self.profitSize))
        print('亏损笔数：{}'.format(self.lossSize))
        print('剩余持仓:{}'.format(self.positionSize))


def main(startcash=1000000, com=0.00016, qts=100):

    # 创建主控制器
    cerebro = bt.Cerebro()
    # 导入策略参数寻优
    cerebro.addstrategy(TestStrategy)
    # cerebro.optstrategy(TestStrategy,maperiod=range(10, 15))
    # 获取数据
    df = writeExcl()  
    df.index = pd.to_datetime(df.trade_time)
    df = df[['open', 'high', 'low', 'close']]
    # 将数据加载至回测系统
    df['SAR1'] = ta.SAR(df.high,df.low,acceleration = 0.1,maximum = 0.2)
    df['SAR2'] = ta.SAR(df.high,df.low,acceleration = 0.05,maximum = 0.05)
    Bup1,Bmd1,Bdn1 = ta.BBANDS(df.SAR1,timeperiod=7*11,matype=5)
    Bup2,Bmd2,Bdn2 = ta.BBANDS(df.SAR2,timeperiod=7*27,matype=5)
    df['Bmd1'] = Bmd1
    df['Bmd2'] = Bmd2
    df['Bdn1'] = Bdn1
    df['Bdn2'] = Bdn2
    print(df)
    df.index = pd.to_datetime( df.index, utc= True)
    data = PandasData_more(
        dataname = df,
        timeframe=bt.TimeFrame.Minutes)
    cerebro.adddata(data)

    # 获取数据
    # query_db = Mysql_search()
    # df = query_db.get_one(['ni'],'2020-06-01','2020-06-03')
    # # ZCL8  JM  J
    # for item in df:
    #     df2 = df[item]
    #     print(df2)
    #     df2['SAR1'] = ta.SAR(df2.high,df2.low,acceleration = 0.1,maximum = 0.2)
    #     df2['SAR2'] = ta.SAR(df2.high,df2.low,acceleration = 0.05,maximum = 0.05)
    #     Bup1,Bmd1,Bdn1 = ta.BBANDS(df2.SAR1,timeperiod=7*11,matype=5)
    #     Bup2,Bmd2,Bdn2 = ta.BBANDS(df2.SAR2,timeperiod=7*27,matype=5)
    #     df2['Bmd1'] = Bmd1
    #     df2['Bmd2'] = Bmd2
    #     df2['Bdn1'] = Bdn1
    #     df2['Bdn2'] = Bdn2
        
    #     df2.index = pd.to_datetime( df2.index, utc= True)
    #     data = PandasData_more(
    #         dataname = df2,
    #         timeframe=bt.TimeFrame.Minutes)
    #     cerebro.adddata(data)


    # broker设置资金、手续费
    cerebro.broker.setcash(startcash)
    cerebro.broker.setcommission(commission=com)
    # 设置买入设置，策略，数量
    cerebro.addsizer(bt.sizers.FixedSize, stake=qts)
    # 以发出信号当日收盘价成交
    cerebro.broker.set_coc(False)

    print('期初总资金: %.2f' % cerebro.broker.getvalue())
    cerebro.addanalyzer(bt.analyzers.TimeReturn, _name='_TimeReturn')

    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name = 'SharpeRatio')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='DW')

    result = cerebro.run()
    print('期末总资金: %.2f' % cerebro.broker.getvalue())
    print('回撤指标:', result[0].analyzers.DW.get_analysis())
    print('夏普比率:', result[0].analyzers.SharpeRatio.get_analysis())

    custom_plot(result)
    # cerebro.plot(volume=False)
    




def writeExcl():
    df = pd.read_csv(r'D:\Downloads\agtd111.csv')
    # df['code'] = 'MAL8'
    return df

main()



